package tjs.study.notes.dotest.proxy.javassist.premain;

import javassist.*;

import java.io.File;
import java.io.FileOutputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;
import java.util.Arrays;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * 字节码转换器
 * 一个提供此接口的实现以转换类文件的代理。转换在main函数之前
 */
public class PremainTransformer implements ClassFileTransformer {

    /**
     * 在main函数执行前，执行的函数（优先agentmain）
     * 使用 premain 可以在类第一次加载之前修改，加载之后修改需要重新创建类加载器。或者在自定义的类加载器种修改
     */
    public static void premain(String options, Instrumentation ins) {
        System.out.println("回调premain.....");
        //注册字节码转换器,有多个，则会调用多次
        ins.addTransformer(new PremainTransformer(), true);//true-允许重新加载class
    }

    /**
     * 字节码加载到虚拟机前会进入这个方法
     * 此方法的实现可以转换提供的类文件，并返回一个新的替换类文件。
     */
    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                            ProtectionDomain protectionDomain, byte[] classfileBuffer) {

        //System.out.println("premain尝试为类： " + className + " 构造代理。。。");
        byte[] res = null;
        //如果加载service包才拦截
        if (className.indexOf("/proxy/service") == -1) {
            return res;
        }

        //javassist的包名是用点分割的，需要转换下
        if (className != null && className.indexOf("/") != -1) {
            className = className.replaceAll("/", ".");
        }
        System.out.println("premain开始为类： " + className + " 构造代理。。。");

        try {
            //通过包名获取类文件
            CtClass ctClass = ClassPool.getDefault().get(className);
            if (ctClass.isFrozen() || ctClass.isInterface()) {
                System.out.println("premain开始为类： " + className + " 已经代理，退出。。。");
                return ctClass.toBytecode();
            }
            //代理目标类的方法
            doProxyMethod(ctClass);

            //生成get/set方法
            doProxyGetSet(ctClass);

            res = ctClass.toBytecode();
            writeClassToDisk(ctClass.getSimpleName(), res);
            return res;
        } catch (Exception ex) {
            ex.printStackTrace();
        }
        return res;
    }

    private void doProxyGetSet(CtClass ctClass) throws CannotCompileException, NotFoundException {
        CtField[] ctFields = ctClass.getDeclaredFields();
        Map<String, CtMethod> allMethodNameMap = Arrays.stream(ctClass.getDeclaredMethods()).collect(Collectors.toMap(CtMethod::getName, Function.identity(), (v1, v2) -> v2));
        for (CtField ctField : ctFields) {
            String paramName = ctField.getName();
            String paramType = ctField.getType().getSimpleName();
            //生成get方法
            CtMethod ctMethodGet = allMethodNameMap.get("getUserName");
            if (ctMethodGet == null) {
                ctMethodGet = CtNewMethod.make(
                        "    public " + paramType + " get" + captureName(paramName) + "() {\n" +
                                "        return " + paramName + " + new java.lang.String(\"make by premain\");\n" +
                                "    }"
                        , ctClass);
                ctClass.addMethod(ctMethodGet);
            } else {
                System.out.println("getUserName已经存在。。。");
            }


            //生成set方法
            CtMethod ctMethodSet = allMethodNameMap.get("setUserName");
            if (ctMethodSet == null) {
                ctMethodSet = CtNewMethod.make(
                        " public void set" + captureName(paramName) + "(" + paramType + " userName) {\n" +
                                "        this." + paramName + " = " + paramName + ";\n" +
                                "    }"
                        , ctClass);
                ctClass.addMethod(ctMethodSet);
            } else {
                System.out.println("setUserName已经存在。。。");
            }

        }
    }

    /**
     * 首字母大写
     *
     * @param name
     * @return
     */
    public static String captureName(String name) {
        char[] cs = name.toCharArray();
        cs[0] -= 32;
        return String.valueOf(cs);

    }

    private void doProxyMethod(CtClass ctClass) throws CannotCompileException {
        CtMethod[] ctMethods = ctClass.getDeclaredMethods();
        for (CtMethod ctMethod : ctMethods) {
            //在方法执行前插入代码
            ctMethod.insertBefore("{ System.out.println(\"前置增强。。。\"); }");
            ctMethod.insertAfter("{ System.out.println(\"后置增强。。。\"); }");
        }
    }

    //将class文件存在磁盘
    private void writeClassToDisk(String className, byte[] data) {
        try {
            String fileName = this.getClass().getResource("").getPath().replaceAll("/target/classes", "/src/main/java");
            fileName = fileName + className + ".class";
            File file = new File(fileName);
            if (!file.exists()) {
                file.createNewFile();
            }
            FileOutputStream fos = new FileOutputStream(file);
            fos.write(data);
            fos.close();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}